﻿using System.Collections.Generic;
using System.Runtime.InteropServices;
using Dou.Job;
using Unity.Burst;
using Unity.Collections;
using Unity.Collections.LowLevel.Unsafe;
using Unity.Jobs;
using Unity.Mathematics;
using UnityEngine;
using UnityEngine.Rendering;
using Random = Unity.Mathematics.Random;

namespace Learning
{
    public class ViewportClipping_CPU : MonoBehaviour
    {
        /// <summary>
        /// GPU Instancing 一个批次可提交的数量
        /// 注意：这里设置得越大，DC会越少，但是特殊情况下浪费的带宽会越多（比如只提交10个时，一个批次还是会提交512个数据给GPU，但是有效数据只有10个）
        /// </summary>
        private const int BATCH = 512;
        
        private const float Range = 100000f;
        private readonly Bounds BOUNDS = new Bounds(Vector3.zero, new Vector3(Range, Range, Range));
        
        public Camera camera;
        public int instanceCount = 100000;
        public Mesh instanceMesh;
        public Material instanceMaterial;
        
        // 所有物体列表
        private NativeArray<VC_Data> _dataArr;
        // 最终筛选出的会进行渲染的列表
        private NativeList<float4x4> _showArr;
        
        private VC_InitJob _initJob;
        private ManagedJobParallelFor<VC_InitJob> _managedIntJob;
        
        private VC_UpdateJob _updateJob;
        private ManagedJob<VC_UpdateJob> _managedUpdateJob;
        private JobHandle _updateHandle;
        
        private List<UnsafeBatchInfo> _pool = new List<UnsafeBatchInfo>();
        
        unsafe void Start()
        {
            _dataArr = new NativeArray<VC_Data>(instanceCount, Allocator.Persistent);
            _showArr = new NativeList<float4x4>(Allocator.Persistent);
            
            _initJob = new VC_InitJob();
            _initJob.dataArr = (VC_Data*) _dataArr.GetUnsafePtr();
            _managedIntJob = new ManagedJobParallelFor<VC_InitJob>();
            _managedIntJob.Work = _initJob;
            
            var handle = _managedIntJob.Schedule(instanceCount, 64);
            handle.Complete();
            _managedIntJob.Dispose();
            _initJob = null;

            _updateJob = new VC_UpdateJob();
            _updateJob.count = instanceCount;
            _updateJob.camera = camera;
            _updateJob.dataArr = (VC_Data*) _dataArr.GetUnsafePtr();
            _updateJob.showArr = _showArr;
            _managedUpdateJob = new ManagedJob<VC_UpdateJob>();
            _managedUpdateJob.Work = _updateJob;
        }

        unsafe void Update()
        {
            FrustumClipping.Update(camera);
            
            _showArr.Clear();
            _updateHandle = _managedUpdateJob.Schedule();
            JobHandle.ScheduleBatchedJobs();
        }

        unsafe void LateUpdate()
        {
            _updateHandle.Complete();

            var batchCount = (int) math.ceil((float) _showArr.Length / (float) BATCH);
            for (int i = 0; i < batchCount; i++)
            {
                var startIndex = i * BATCH;
                var len = BATCH;
                if (i == batchCount - 1)
                {
                    len = _showArr.Length - (BATCH * i);
                }
                
                if (_pool.Count == i)
                {
                    var newInfo = new UnsafeBatchInfo();
                    
                    newInfo.batchTransArray = new NativeArray<float4x4>(BATCH, Allocator.Persistent);
                    newInfo.transBuf = new GraphicsBuffer(GraphicsBuffer.Target.Structured, BATCH, Marshal.SizeOf(typeof(float4x4)));
                    newInfo.transBuf.SetData(newInfo.batchTransArray);
                    
                    newInfo.block = new MaterialPropertyBlock();
                    newInfo.block.SetBuffer("transBuf", newInfo.transBuf);
                    
                    newInfo.args = new uint[5] {0, 0, 0, 0, 0};
                    newInfo.args[0] = instanceMesh.GetIndexCount(0); // 每个实例的索引数
                    newInfo.args[1] = 0; // 实例数
                    newInfo.args[2] = instanceMesh.GetIndexStart(0); // 起始索引位置
                    newInfo.args[3] = instanceMesh.GetBaseVertex(0); // 基顶点位置
                    newInfo.args[4] = 0; // 起始实例位置
                    newInfo.argsBuf = new GraphicsBuffer(GraphicsBuffer.Target.IndirectArguments, 1, 5 * sizeof(uint));
                    newInfo.argsBuf.SetData(newInfo.args);
                    
                    _pool.Add(newInfo);
                }
                
                var info = _pool[i];
                
                // 填充数据
                var sourcePtr = (float4x4*) _showArr.GetUnsafePtr();
                sourcePtr += startIndex;
                var destPtr = (float4x4*) info.batchTransArray.GetUnsafePtr();
                UnsafeUtility.MemCpy(destPtr, sourcePtr, BATCH * 16 * 4); // 直接通过指针拷贝数据
                info.transBuf.SetData(info.batchTransArray); // 数据变动后需要重新设置才能生效
                
                info.args[1] = (uint) len; // 实例数
                info.argsBuf.SetData(info.args); // 数据变动后需要重新设置才能生效
                
                // 绘制
                Graphics.DrawMeshInstancedIndirect(instanceMesh, 0, instanceMaterial, BOUNDS, info.argsBuf,
                    0, info.block, ShadowCastingMode.On, true, LayerMask.NameToLayer("Character"));
            }
        }

        void OnDestroy()
        {
            _managedUpdateJob.Dispose();
            _updateJob = null;
        }
    }

    /// <summary>
    /// 单个物体的数据，同一个Mesh和Mat渲染的物体，会有一个VC_Data[]记录，这里我们只使用一个Mesh和Mat
    /// </summary>
    public struct VC_Data
    {
        public float4x4 trans;
        public float3 center;
        public float3 extents;
    }

    /// <summary>
    /// 初始化所有物体
    /// </summary>
    [BurstCompile]
    public unsafe class VC_InitJob : ManagedJobParallelFor<VC_InitJob>.IWork
    {
        public VC_Data* dataArr;
        
        [BurstCompile]
        [Il2CppSetOption(Option.ArrayBoundsChecks, false)]
        [Il2CppSetOption(Option.NullChecks, false)]
        public void Execute(int index)
        {
            var random = Random.CreateFromIndex((uint) index);
            
            float angle = random.NextFloat(0.0f, Mathf.PI * 2.0f);
            float distance = random.NextFloat(20.0f, 100.0f);
            float height = random.NextFloat(-2.0f, 2.0f);
            float3 pos = new float3(math.sin(angle) * distance, height, math.cos(angle) * distance);
            
            quaternion rotation = quaternion.Euler(random.NextFloat3(0, 360));
            float3 scale = random.NextFloat3(0.05f, 0.25f);
            
            var data = new VC_Data();
            data.trans = float4x4.TRS(pos, rotation, scale);
            data.center = pos;
            data.extents = scale * 0.5f; // 包围盒一半的尺寸
            
            var ptr = dataArr + index;
            *ptr = data;
        }
    }
    
    /// <summary>
    /// 每帧执行视椎剔除
    /// </summary>
    [BurstCompile]
    public unsafe class VC_UpdateJob : ManagedJob<VC_UpdateJob>.IWork
    {
        public Camera camera;
        
        public float count;
        public VC_Data* dataArr;
        public NativeList<float4x4> showArr;
        
        [BurstCompile]
        [Il2CppSetOption(Option.ArrayBoundsChecks, false)]
        [Il2CppSetOption(Option.NullChecks, false)]
        public void Execute()
        {
            for (int i = 0; i < count; i++)
            {
                var data = *(dataArr + i);
                // 检测是否位于视椎体中
                if (FrustumClipping.GetClipping(data.center, data.extents) != ClippingType.Out)
                {
                    showArr.Add(data.trans); // Add 方法线程不安全，如果是使用 IJobParallelFor 则这里需要加锁，否则 Unity 会奔溃
                }
            }
        }
    }
}
